import pandas as pd
import numpy as np
from datetime import datetime
from sqlalchemy import create_engine, text
from sqlalchemy.engine import URL
from urllib.parse import quote_plus
import sys
import traceback
from numba import jit, prange
import re

# 初始化全局计数器
out_of_range_total = 0
abnormal_age_data = []  # 用于存储异常年龄数据

# 定义非法字符的正则表达式模式
illegal_pattern = re.compile(r'[\x00-\x08\x0B\x0C\x0E-\x1F\x7F-\x9F]')

def remove_illegal_characters(text):
    """移除字符串中的非法字符"""
    if isinstance(text, str):
        return illegal_pattern.sub('', text)
    return text

@jit(nopython=True, parallel=True)
def calculate_age_numba(birth_dates, today_year, today_month, today_day):
    """使用Numba JIT加速的年龄计算函数（并行计算）"""
    ages = np.zeros(len(birth_dates), dtype=np.int32)
    for i in prange(len(birth_dates)):
        year, month, day = birth_dates[i]
        if year == 0:  # 无效日期
            ages[i] = -1
            continue
            
        age = today_year - year
        if (month, day) > (today_month, today_day):
            age -= 1
            
        ages[i] = age
    return ages

def get_db_engine(database_config):
    """获取优化后的数据库引擎"""
    db_type = database_config.get('type', 'postgresql')
    
    if db_type == 'postgresql':
        return create_engine(
            f"postgresql+psycopg2://{database_config['user']}:{quote_plus(database_config['password'])}@"
            f"{database_config['host']}:{database_config['port']}/{database_config['database']}",
            pool_size=30,
            max_overflow=100,
            pool_pre_ping=True
        )
    elif db_type == 'mysql':
        return create_engine(
            f"mysql+pymysql://{database_config['user']}:{quote_plus(database_config['password'])}@"
            f"{database_config['host']}:{database_config['port']}/{database_config['database']}",
            pool_size=30,
            max_overflow=100
        )
    else:
        raise ValueError(f"不支持的数据库类型: {db_type}")

def optimized_age_calc(df):
    """从证件号码提取出生年月日并计算年龄（超高性能实现）"""
    global out_of_range_total
    
    # 初始化出生日期列为空
    df['出生日期'] = pd.NaT
    
    # 提取身份证号码中的出生年月日
    id_number = df['证件号码'].astype(str).str.strip()
    
    # 处理18位身份证号码（第7-14位为出生日期）
    is_18 = id_number.str.len() == 18
    if is_18.any():
        birth_18 = id_number[is_18].str[6:14]
        df.loc[is_18, '出生日期'] = pd.to_datetime(birth_18, format='%Y%m%d', errors='coerce')
    
    # 转换出生日期为年、月、日数组
    birth_dates = df['出生日期'].dt
    birth_years = birth_dates.year.fillna(0).astype(np.int32)
    birth_months = birth_dates.month.fillna(1).astype(np.int32)
    birth_days = birth_dates.day.fillna(1).astype(np.int32)
    
    # 今天的日期（2025-01-01）
    today = datetime(2025, 1, 1)
    
    # 使用Numba JIT并行计算年龄
    ages = calculate_age_numba(
        np.column_stack((birth_years, birth_months, birth_days)),
        today.year, today.month, today.day
    )
    
    # 转换为Pandas Series
    df['计算年龄'] = pd.Series(ages, index=df.index)
    
    # 统计超出范围的年龄（0-120岁）
    out_of_range = (df['计算年龄'] < 6) | (df['计算年龄'] > 120) | (df['计算年龄'] == -1)
    out_of_range_total += out_of_range.sum()
    
    # 标记显示年龄（根据身份证号码长度处理）
    df['显示年龄'] = df.apply(lambda row: 
                             row['计算年龄'] if (row['证件号码'] and len(str(row['证件号码'])) == 18) 
                             else '', axis=1)
    
    return df

def get_service_stats(engine):
    """获取服务时长统计数据"""
    print("获取服务时长统计数据...")
    with engine.connect() as conn:
        # 2025年服务时长>0的用户数
        query_2025 = """
            SELECT axcp0005 FROM vsis.kdx07 
            WHERE albx0702 IS NOT NULL AND albx0702::numeric > 0.0  
            AND (substring(albx0703, 1, 4) = '2025' 
            OR substring(albx0704, 1, 4) = '2025' 
            OR substring(update_time::text, 1, 4) = '2025')
            GROUP BY axcp0005
        """
        result_2025 = conn.execute(text(query_2025))
        users_2025 = {row[0] for row in result_2025}
        
        # 往年服务时长>0的用户数
        query_previous = """
            SELECT axcp0005 FROM vsis.kdx07 
            WHERE albx0702 IS NOT NULL AND albx0702::numeric > 0.0  
            AND (substring(albx0703, 1, 4) != '2025' 
            AND substring(albx0704, 1, 4) != '2025' 
            AND substring(update_time::text, 1, 4) != '2025')
            GROUP BY axcp0005
        """
        result_previous = conn.execute(text(query_previous))
        users_previous = {row[0] for row in result_previous}
        
    return users_2025, users_previous

def analyze_data(database_config, output_file=None, chunksize=2000000):
    """超高性能数据分析函数（基于证件号码计算年龄）"""
    try:
        # 获取数据库引擎
        engine = get_db_engine(database_config)
        
        # 预先获取服务统计数据
        users_2025, users_previous = get_service_stats(engine)
        
        # 执行综合统计查询
        print("执行综合统计查询...")
        with engine.connect() as conn:
            combined_query = """
                WITH stats AS (
                    SELECT 
                        COUNT(CASE WHEN axcp0077='0' THEN 1 END) as 注册人数,
                        SUM(CASE WHEN axcp0077='0' AND axcp0004='01' AND LENGTH(axcp0005)=18 THEN 1 ELSE 0 END) as 有效身份证件用户数,
                        SUM(CASE WHEN axcp0077='0' AND axcp0004!='01' THEN 1 ELSE 0 END) as 非身份证件用户数,
                        SUM(CASE WHEN (albp0289='0.0' OR albp0289 IS NULL) AND axcp0077='0' THEN 1 ELSE 0 END) as 志愿时长为0的用户数,
                        SUM(CASE WHEN axcp0077='0' AND axcp0004='01' AND (axcp0005 is null OR LENGTH(axcp0005)!=18) THEN 1 ELSE 0 END) as 身份证缺失人数
                    FROM vsis.kdp01
                ),
                duplicate_ids AS (
                    SELECT COUNT(1) as 身份证重复人数
                    FROM (
                        SELECT axcp0005 
                        FROM vsis.kdp01 
                        WHERE axcp0004='01' AND axcp0005!='' AND axcp0077='0'
                        GROUP BY axcp0005 
                        HAVING COUNT(1)>1
                    ) t
                )
                SELECT * FROM stats, duplicate_ids
            """
            result = conn.execute(text(combined_query))
            basic_stats = dict(zip(result.keys(), result.fetchone()))
            
            # 添加服务统计数据
            basic_stats['2025年服务时长>0的用户数'] = len(users_2025)
            basic_stats['往年服务时长>0的用户数'] = len(users_previous)
            print("基本统计结果:", basic_stats)
        
        # 分块读取数据（新增姓名、登记日期等字段）
        print(f"开始分块读取kdp01表数据，块大小: {chunksize}")
        columns = ['axcp0039', 'axcp0004', 'axcp0005', 'albp0289', 'axcp0077', 'axcp0002', 'axbe0031']
        query = f"SELECT {', '.join(columns)} FROM vsis.kdp01"
        
        age_group_stats = pd.DataFrame()
        total_records = 0
        valid_ages_total = 0
        start_time = datetime.now()
        
        # 分块处理数据
        for i, chunk in enumerate(pd.read_sql(query, engine, chunksize=chunksize)):
            print(f"处理数据块 {i+1} ({len(chunk)} 条记录)")
            total_records += len(chunk)
            
            # 重命名列
            df = chunk.rename(columns={
                'axcp0005': '证件号码',
                'axcp0004': '证件类型',
                'albp0289': '服务时长',
                'axcp0077': '注册状态',
                'axcp0002': '姓名',
                'axbe0031': '登记日期'
            })
            
            # 数据预处理
            df['服务时长'] = pd.to_numeric(df['服务时长'], errors='coerce').fillna(0).astype(np.float32)
            
            # 从证件号码计算年龄
            df = optimized_age_calc(df)
            valid_ages = (~df['计算年龄'].isnull() & (df['计算年龄'] != -1)).sum()
            valid_ages_total += valid_ages
            
            # 年龄分组
            age_bins = [0, 6, 10, 18, 30, 40, 50, 60, 65, 70, 80, 90, 100, 120, float('inf')]
            age_labels = ['≤6', '6＜年龄≤10', '10＜年龄≤18', '18＜年龄≤30', 
                         '30＜年龄≤40', '40＜年龄≤50', '50＜年龄≤60', 
                         '60＜年龄≤65', '65＜年龄≤70', '70＜年龄≤80', 
                         '80＜年龄≤90', '90＜年龄≤100', '100＜年龄≤120', '＞120']
            df['年龄段'] = pd.cut(df['计算年龄'], bins=age_bins, labels=age_labels, right=False)
            df['年龄段'] = df['年龄段'].cat.add_categories(['未知年龄']).fillna('未知年龄')
            
            # 标记证件状态
            df['有效身份证件'] = (df['证件类型'] == '01') & (df['证件号码'].str.len() == 18)
            df['证件号缺失'] = (df['证件类型'] == '01') & ~df['有效身份证件']
            df['非身份证件'] = df['证件类型'] != '01'
            df['服务时长为0'] = (df['服务时长'] == 0) | df['服务时长'].isna()
            
            # 标记服务状态
            df['2025年服务时长>0'] = df['证件号码'].isin(users_2025)
            df['往年服务时长>0'] = df['证件号码'].isin(users_previous)
            
            # 筛选异常年龄数据（小于6岁、大于120岁和未知年龄）
            abnormal_age_mask = (df['年龄段'] == '≤6') | (df['年龄段'] == '＞120') | (df['年龄段'] == '未知年龄')
            if abnormal_age_mask.any():
                abnormal_chunk = df[abnormal_age_mask].copy()
                # 只保留需要的字段
                abnormal_chunk = abnormal_chunk[['姓名', '证件号码', '登记日期', '显示年龄', '年龄段']]
                
                # 清洗数据中的非法字符
                for col in ['姓名', '证件号码', '登记日期']:
                    abnormal_chunk[col] = abnormal_chunk[col].apply(remove_illegal_characters)
                
                abnormal_age_data.append(abnormal_chunk)
            
            # 分组统计
            registered_df = df[df['注册状态'] == '0'].copy()
            if not registered_df.empty:
                groups = registered_df['年龄段'].unique()
                chunk_stats = []
                for group in groups:
                    group_df = registered_df[registered_df['年龄段'] == group]
                    chunk_stats.append([
                        group,
                        len(group_df),
                        group_df['证件号码'].duplicated().sum(),
                        group_df['证件号缺失'].sum(),
                        group_df['有效身份证件'].sum(),
                        group_df['非身份证件'].sum(),
                        group_df['服务时长为0'].sum(),
                        group_df['2025年服务时长>0'].sum(),
                        group_df['往年服务时长>0'].sum()
                    ])
                if chunk_stats:
                    chunk_stats_df = pd.DataFrame(
                        chunk_stats,
                        columns=['年龄段', '年龄段计数', '证件号码重复数', 
                                 '证件号缺失数', '有效身份证件用户数', 
                                 '非身份证件用户数', '服务时长为0的用户数',
                                 '2025年服务时长>0的用户数', '往年服务时长>0的用户数']
                    )
                    age_group_stats = pd.concat([age_group_stats, chunk_stats_df], ignore_index=True)
            
            # 进度提示
            elapsed = (datetime.now() - start_time).total_seconds()
            if (i+1) % 10 == 0 or i == 0:
                print(f"已处理: {total_records:,} 条, 耗时: {elapsed:.2f}秒")
        
        # 合并统计结果
        if not age_group_stats.empty:
            age_group_stats = age_group_stats.groupby('年龄段').sum().reset_index()
        
        # 合并异常年龄数据
        if abnormal_age_data:
            abnormal_age_df = pd.concat(abnormal_age_data, ignore_index=True)
            # 重命名列以便清晰显示
            abnormal_age_df = abnormal_age_df.rename(columns={
                '姓名': 'axcp0002 姓名',
                '证件号码': 'axcp0005 人员身份证件号码',
                '登记日期': 'axbe0031 登记日期',
                '显示年龄': '年龄（从身份证号码中截取）',
                '年龄段': '异常类型'
            })
            
            # 再次清洗最终结果中的非法字符
            for col in abnormal_age_df.columns:
                if col != '年龄（从身份证号码中截取）':  # 跳过数字列
                    abnormal_age_df[col] = abnormal_age_df[col].apply(remove_illegal_characters)
            
            print(f"发现 {len(abnormal_age_df)} 条异常年龄数据")
        else:
            abnormal_age_df = pd.DataFrame()
            print("未发现异常年龄数据")
        
        # 输出结果
        print(f"\n总记录数: {total_records:,}, 成功计算年龄: {valid_ages_total:,}")
        print(f"年龄超出范围总数: {out_of_range_total:,}")
        
        # 保存结果
        if output_file:
            print(f"保存结果到: {output_file}")
            with pd.ExcelWriter(output_file) as writer:
                pd.DataFrame(basic_stats.items(), columns=['统计项', '数量']).to_excel(
                    writer, sheet_name='基本统计', index=False)
                age_group_stats.to_excel(writer, sheet_name='年龄分组统计', index=False)
                
                # 写入异常年龄数据
                if not abnormal_age_df.empty:
                    abnormal_age_df.to_excel(writer, sheet_name='异常年龄数据', index=False)
        
        end_time = datetime.now()
        print(f"总耗时: {(end_time - start_time).total_seconds():.2f}秒")
        return {
            'basic_stats': basic_stats,
            'age_group_stats': age_group_stats,
            'abnormal_age_data': abnormal_age_df,
            'total_time': (end_time - start_time).total_seconds()
        }
        
    except Exception as e:
        print(f"分析错误: {e}")
        traceback.print_exc()
        return None
    finally:
        engine.dispose()

if __name__ == "__main__":
    # PostgreSQL数据库配置（请根据实际情况修改）
    database_config = {
        'type': 'postgresql',
        'host': '192.1.0.33',
        'user': 'vsisuser',
        'password': '123Qwe!@#',
        'database': 'vsisdb',
        'port': 5432,
        'schema': 'vsis'
    }

    # 输出文件路径
    output_file = r'/database/csv_file/script/result.xlsx'

    # 执行分析
    print("开始分析数据...")
    results = analyze_data(
        database_config, 
        output_file=output_file,
        chunksize=2000000  # 每块200万条数据，根据内存调整
    )
    
    if results:
        print(f"数据分析完成! 总耗时: {results['total_time']:.2f}秒")
        if not results['abnormal_age_data'].empty:
            print(f"已将 {len(results['abnormal_age_data'])} 条异常年龄数据写入'异常年龄数据'工作表")